#####   SOLVER FOR STAGES 1 and 2 #####
# -------------------------------------------------------------------------------
function solve_main(ce::CE_Economy_PInf, EWf_ΠT::Array{Float64,3})
  #  State: (b,y,j); j=1 is P; j=2 is I
# -------------------------------------------------------------------------------
  
#Create matrices to fill-in
  bp_mat  = Array{Float64}(ce.nb, ce.ny, 2);
  WR_mat  = zeros(ce.nb, ce.ny, 2);

#Create interpolation objects
  knots   = (ce.b_grid, ce.y_grid, 1:2);
  itp_WΠT = interpolate(knots, EWf_ΠT,     Gridded(Linear()));
  itp_q   = interpolate(knots, ce.q_mat,   Gridded(Linear()));

for i_j               = 1:2
for (i_y, y_v)        in enumerate(ce.y_grid)
for (i_b, b_v)        in enumerate(ce.b_grid)
  πv = ifelse(i_j==1, 0.0, ce.π_min)


    # ----------------------------------------------------------------------------
    # SOLVE FOR OPTIMAL BOND POLICIES
    # ----------------------------------------------------------------------------
    fun_bp_fx(bp_v::Float64) =  fun_bp(ce, bp_v, πv, b_v, y_v, ce.Bval, itp_q[bp_v,y_v,i_j], itp_WΠT[bp_v,y_v,i_j])
    #Solver
    res_bp   =  optimize(fun_bp_fx, ce.b_grid[1], ce.b_grid[end], Brent(),  abs_tol=1e-5, iterations=300)
    bp_v     =  Optim.minimizer(res_bp)
    val_max  =  -Optim.minimum(res_bp)
        #Compare with lower bound solution
          tmp_val   = -fun_bp_fx(ce.b_grid[1])
          if tmp_val > val_max
          val_max   = tmp_val
          bp_v      = ce.b_grid[1]
          end
        #Check previous solution
          bp_prev_sol =ce.bp[i_b,i_y,i_j]
          tmp_val =  -fun_bp_fx(bp_prev_sol)
          if tmp_val > val_max
          val_max   = tmp_val
          bp_v      = bp_prev_sol
          end
      #Store solutions
    WR_mat[i_b,i_y,i_j] =  tmp_val
    bp_mat[i_b,i_y,i_j] =  bp_v;
end # b
end # y
end # j

#STORE RESULTS:
  dist_b_abs = maximum(abs, ce.bp-bp_mat)
  copy!(ce.bp, bp_mat)
  copy!(ce.Wr, WR_mat);
return dist_b_abs
end
# -------------------------------------------------------------------------------


# -------------------------------------------------------------------------------
# Bond pricing kernel, q(.)
# -------------------------------------------------------------------------------
function compute_prices(ce::CE_Economy_PInf)

  # Create matrix to fill in
  q_new       = Array{Float64,3}(size(ce.q_mat)) 
  # Interpolating objects
  knots       = (ce.b_grid, ce.y_grid, 1:2)
  default_int = interpolate(knots,  ce.default,    Gridded(Linear()))
  bp_int      = interpolate(knots,  ce.bp,         Gridded(Linear()))
  # Interpolate relevant prices
  q_int       = interpolate(knots,  ce.q_mat,      Gridded(Linear()))

  for i_j=1:2
  for (i_y, y_v)   in enumerate(ce.y_grid)
  for (i_bp, bp_v) in enumerate(ce.b_grid)
      #Initialize recursion
      q_v = 0.0;
      for (i_yp, yp_v)   in enumerate(ce.y_grid)
      for i_jp=1:2
          #Next period default policy
          df_bel = default_int[bp_v, yp_v, i_jp]
          #Next-period bond policy
          bpp    = bp_int[bp_v, yp_v, i_jp]
          #Next period price
          qp = q_int[bpp,yp_v,i_jp]
          #Update recursion
          q_v += ce.Πy[i_y, i_yp] * ce.ΠT[i_j, i_jp] * (1.0-df_bel) * ((1-ce.mb)*(ce.zb+qp )+ce.mb);
      end #close jp
      end #close yp
      q_new[i_bp, i_y, i_j]    =  q_v / (1.0+ce.r)
  end
  end
  end

  #Compute distances
  dist_q  	         = norm(vec( q_new - ce.q_mat));
  dist_q_abs         = maximum(abs, q_new - ce.q_mat);
  copy!(ce.q_mat, ce.q_mat*ce.damp + q_new*(1-ce.damp))
return dist_q, dist_q_abs
end
# -----------------------------------------------------------------------------

#####   STAGE 0 #####
# -------------------------------------------------------------------------
function compute_def_prob(ce::CE_Economy_PInf, Wr::Float64, Wd::Float64)
  μ        = 1.0
  σ        = ce.σ_d
  a        = Wr/Wd
  prob_def = cdf(Normal(μ,σ), a)
return prob_def
end
# --------------------------------------------------------------------------

# -------------------------------------------------------------------------------
function solve_default(ce::CE_Economy_PInf, EWf::Array{Float64,3}, EWd::Array{Float64,2})
 
#Create matrices to store results
 Wf_new   = Array{Float64}(ce.nb, ce.ny, 2);
 Wd_new   = Array{Float64}(ce.ny, 2);
 default  = Array{Float64}(ce.nb, ce.ny, 2);
#Define interpolation objects
 knots    = (ce.b_grid, ce.y_grid, 1:2)
 knots_def= (ce.y_grid, 1:2)
 EWf_int  = interpolate(knots,     EWf,    Gridded(Linear()));
 Wr_int   = interpolate(knots,     ce.Wr,  Gridded(Linear()));
 EWd_int  = interpolate(knots_def, EWd,    Gridded(Linear()));

for i_j=1:2
for (i_y, y_v) in enumerate(ce.y_grid)
for (i_b, b_v) in enumerate(ce.b_grid)

        Wr_val           = Wr_int[b_v, y_v, i_j]
        Wd_new[i_y, i_j] = ce.Udef[i_y, i_j] + ce.β*( ce.θ*EWf_int[0., y_v, i_j] + (1.0-ce.θ)* EWd_int[y_v, i_j] )

  # Randomizing the default decision
  # -------------------------------------------------------------------------
  DefProb                = compute_def_prob(ce, Wr_val, Wd_new[i_y, i_j] );
  default[i_b, i_y, i_j] = DefProb
  #Value Function (fx of b,y,j)
  Wf_new[i_b, i_y, i_j]  =  DefProb * Wd_new[i_y, i_j] + (1-DefProb) * Wr_val;

end
end
end

#Compute Distances and Update
# ------------------------------------------------------------------------
#Distance between value fx across iterations
  dist_Wf  	   = norm(vec(    Wf_new  - ce.Wf))
  dist_Wf_abs  = maximum(abs, Wf_new  - ce.Wf )
  dist_def     = norm(vec(    default  - ce.default))

#Store results
  copy!(ce.Wf, Wf_new)
  copy!(ce.Wd, Wd_new)
  copy!(ce.default, default)
return dist_Wf, dist_Wf_abs, dist_def
end
# -------------------------------------------------------------------------


# -------------------------------------------------------------------------------
# Computing Expectations of v [includes Markov switches]
# -------------------------------------------------------------------------------
function compute_EV(ce::CE_Economy_PInf, W::Array{Float64})
  EW = Array{Float64}(size(W))
  if length(size(W))==3
  #Case in which government repays [states: y,b,type]
    for i_j=1:2
      for (i_y, y_v)     in enumerate(ce.y_grid)
        for (i_bp, bp_v) in enumerate(ce.b_grid)
            EV_tmp = 0.0
              for i_yp = 1:ce.ny
              for i_jp = 1:2
                EV_tmp += ce.Πy[i_y, i_yp] * ce.ΠT[i_j, i_jp] * W[i_bp, i_yp, i_jp]
              end
              end
            EW[i_bp, i_y, i_j] = EV_tmp
        end
      end
    end
  #Case in which government defaults [states: y,type]
  else
    for i_j=1:2
      for (i_y, y_v)     in enumerate(ce.y_grid)
        EV_tmp = 0.0
          for i_yp = 1:ce.ny
          for i_jp = 1:2
            EV_tmp += ce.Πy[i_y, i_yp] * ce.ΠT[i_j, i_jp] * W[i_yp, i_jp]
          end
          end
        EW[i_y, i_j] = EV_tmp
      end
    end
  end

return EW
end
# -------------------------------------------------------------------------------
